
/******************************************************************************
 *
 *  This file is part of canu, a software program that assembles whole-genome
 *  sequencing reads into contigs.
 *
 *  This software is based on:
 *    'Celera Assembler' r4587 (http://wgs-assembler.sourceforge.net)
 *    the 'kmer package' r1994 (http://kmer.sourceforge.net)
 *
 *  Except as indicated otherwise, this is a 'United States Government Work',
 *  and is released in the public domain.
 *
 *  File 'README.licenses' in the root directory of this distribution
 *  contains full conditions and disclaimers.
 */

#include "stashContains.H"

//  Replace the children list in tig with one that has fewer contains.  The original
//  list is returned.
savedChildren *
stashContains(tgTig       *tig,
              double       maxCov,
              bool         beVerbose) {

  if (tig->numberOfChildren() == 1)
    return(NULL);

  //  Stats we report
  int32  nOrig     = tig->numberOfChildren();
  int32  nBack     = 0;
  int32  nCont     = 0;
  int32  nSave     = 0;
  int64  nBase     = 0;
  int64  nBaseDove = 0;
  int64  nBaseCont = 0;

  //  Sort the original children by position.

  std::sort(tig->_children, tig->_children + tig->_childrenLen);

  //  Decide which children to save.

  bool            *isBack = new bool       [nOrig];   //  True, we save the child for processing
  readLength      *posLen = new readLength [nOrig];   //  Sorting by length of child

  //  The first read is always saved

  int32         loEnd = tig->_children[0].min();
  int32         hiEnd = tig->_children[0].max();

  isBack[0]      = 1;
  nBack          = 1;
  posLen[0].idx  = 0;
  posLen[0].len  = hiEnd - loEnd;
  nBaseDove     += posLen[0].len;
  nBase         += posLen[0].len;

  //  For the other reads, save it if it extends the backbone sequence.

  for (uint32 fi=1; fi<nOrig; fi++) {
    int32  lo = tig->_children[fi].min();
    int32  hi = tig->_children[fi].max();

    posLen[fi].idx  = fi;
    posLen[fi].len  = hi - lo;
    nBase          += posLen[fi].len;

    if (hi <= hiEnd) {
      isBack[fi] = false;
      nCont++;
      nBaseCont += posLen[fi].len;

    } else {
      isBack[fi] = true;
      nBack++;
      nBaseDove += posLen[fi].len;
    }

    hiEnd = max(hi, hiEnd);
  }

  //  Throw out some of the contained reads to make our coverage acceptable.

  std::sort(posLen, posLen + nOrig, greater<readLength>());  //  Sort by length, larger first

  for (uint32 ii=1; ii<nOrig; ii++)                 //  Ensure we're sorted.
    assert(posLen[ii-1].len >= posLen[ii].len);

  int64  saveLimit = maxCov * (int64)hiEnd - nBaseDove;
  int64  nBaseSave = 0;

  for (uint32 ii=0; ii<nOrig; ii++) {
    if (isBack[posLen[ii].idx] == true)    //  Already a backbone read.
      continue;                            //  Skip this read.

    if (nBaseSave > saveLimit)             //  Exceeded coverage limit.
      break;                               //  Bail.

    isBack[posLen[ii].idx] = true;
    nSave++;
    nBaseSave += posLen[ii].len;
  }

  //  Initialize the savedChuldren statistics.

  savedChildren   *saved = new savedChildren();

  saved->numContains  = nCont;
  saved->covContain   = (double)nBaseCont / hiEnd;
  saved->percContain  = 100.0 * nBaseCont / nBase;;

  saved->numDovetails = nBack;
  saved->covDovetail  = (double)nBaseDove / hiEnd;
  saved->percDovetail = 100.0 * nBaseDove / nBase;;

  saved->numContainsSaved   = nSave;
  saved->covContainsSaved   = (double)nBaseSave / hiEnd;

  saved->numContainsRemoved = nOrig - nBack - nSave;
  saved->covContainsRemoved = (double)(nBaseCont - nBaseSave) / hiEnd;

  //  If we've flagged stuff for removal, remove them.  Otherwise, coverage
  //  is acceptable and we didn't do anything to the list of children
  //  (except sort by position).

  if (saved->numContainsRemoved > 0) {
    saved->childrenLen = tig->_childrenLen;
    saved->childrenMax = tig->_childrenMax;
    saved->children    = tig->_children;

    tig->_childrenLen  = 0;
    tig->_childrenMax  = nBack + nSave;
    tig->_children     = new tgPosition [tig->_childrenMax];

    for (uint32 fi=0; fi<nOrig; fi++)
      if (isBack[fi] == true)
        tig->_children[tig->_childrenLen++] = saved->children[fi];
  }

  //  Cleanup and return the saved children (or an empty list if nothing was saved).

  delete [] isBack;
  delete [] posLen;

  return(saved);
}


//  Restores the f_list, and updates the position of non-contained reads.
//
void
unstashContains(tgTig                *tig,
                savedChildren        *saved) {

  if ((saved == NULL) ||
      (saved->numContainsRemoved == 0))
    return;

  //  For fragments not involved in the consensus computation, we'll scale their position linearly
  //  from the old max to the new max.
  //
  //  We probably should do an alignment to the consensus sequence to find the true location, but
  //  that's (a) expensive and (b) likely overkill for these unitigs.

  uint32   oldMax = 0;
  uint32   newMax = 0;
  double   sf     = 1.0;

  for (uint32 fi=0, ci=0; fi<saved->childrenLen; fi++)
    if (oldMax < saved->children[fi].max())
      oldMax = saved->children[fi].max();

  for (uint32 fi=0, ci=0; fi<tig->numberOfChildren(); fi++)
    if (newMax < tig->getChild(fi)->max())
      newMax = tig->getChild(fi)->max();

  if (oldMax > 0)
    sf = (double)newMax / oldMax;

  //  First, we need a map from the child id to the location in the current tig

  map<int32, tgPosition *>   idmap;

  for (uint32 ci=0; ci < tig->numberOfChildren(); ci++)
    idmap[tig->getChild(ci)->ident()] = tig->getChild(ci);

  //  Now, over all the reads in the original saved fragment list, update the position.  Either from
  //  the computed result, or by extrapolating.

  for (uint32 fi=0; fi<saved->childrenLen; fi++) {
    uint32  iid = saved->children[fi].ident();

    //  Does the ID exist in the new positions?  Copy the new position to the original list.
    if (idmap.find(iid) != idmap.end()) {
      saved->children[fi] = *idmap[iid];
      idmap.erase(iid);
    }

    //  Otherwise, fudge the positions.
    else {
      int32  nmin = sf * saved->children[fi].min();
      int32  nmax = sf * saved->children[fi].max();

      if (nmin > newMax)  nmin = newMax;
      if (nmax > newMax)  nmax = newMax;

      saved->children[fi].setMinMax(nmin, nmax);
    }
  }

  if (idmap.empty() == false)
    fprintf(stderr, "Failed to unstash the contained reads.  Still have " F_SIZE_T " reads unplaced.\n",
            idmap.size());
  assert(idmap.empty() == true);

  //  Throw out the reduced list, and restore the original.

  delete [] tig->_children;

  tig->_childrenLen  = saved->childrenLen;
  tig->_childrenMax  = saved->childrenMax;
  tig->_children     = saved->children;

  saved->childrenLen = 0;
  saved->childrenMax = 0;
  saved->children    = NULL;
}

